{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure AI Search with LlamaIndex\n",
    "\n",
    "This code demonstrates how to use Azure AI Search with Azure OpenAI and the [LlamaIndex data framework](https://www.llamaindex.ai/).\n",
    "\n",
    "You need an existing [Azure AI Search service](https://learn.microsoft.com/azure/search/search-create-service-portal), any tier, but it must have capacity and quota for the workload. We recommend Basic or higher for this sample.\n",
    "\n",
    "You also need [Azure OpenAI](https://learn.microsoft.com/azure/ai-services/openai/how-to/create-resource), with deployments of a text embedding model such as [text-embedding-ada-002](https://learn.microsoft.com/azure/ai-services/openai/concepts/models#embeddings), and a chat model such as [gpt-35-turbo](https://learn.microsoft.com/azure/ai-services/openai/concepts/models#gpt-35)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up a Python virtual environment in Visual Studio Code\n",
    "\n",
    "1. Open the Command Palette (Ctrl+Shift+P).\n",
    "1. Search for **Python: Create Environment**.\n",
    "1. Select **Venv**.\n",
    "1. Select a Python interpreter. Choose 3.10 or later.\n",
    "\n",
    "It can take a minute to set up. If you run into problems, see [Python environments in VS Code](https://code.visualstudio.com/docs/python/environments)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -r azure-search-vector-python-llamaindex-sample-requirements.txt --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load .env file\n",
    "\n",
    "Copy `/code/.env-sample` to an `.env` file in the sample folder, and update accordingly. The search service and Azure OpenAI resource and models must exist, but the search index is created and loaded during code execution. Provide a unique name for the search index. Endpoints, API keys, and Azure OpenAI information can be found in the Azure portal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "import os\n",
    "\n",
    "load_dotenv(override=True) # take environment variables from .env.\n",
    "\n",
    "# Make sure your .env file has values for the following environment variables\n",
    "endpoint = os.environ[\"AZURE_SEARCH_SERVICE_ENDPOINT\"]\n",
    "credential = AzureKeyCredential(os.environ[\"AZURE_SEARCH_ADMIN_KEY\"]) if len(os.environ[\"AZURE_SEARCH_ADMIN_KEY\"]) > 0 else DefaultAzureCredential()\n",
    "index_name = os.environ[\"AZURE_SEARCH_INDEX\"]\n",
    "azure_openai_endpoint = os.environ[\"AZURE_OPENAI_ENDPOINT\"]\n",
    "# Llama Index does not support RBAC authentication, an API key is required\n",
    "azure_openai_key = os.environ[\"AZURE_OPENAI_KEY\"]\n",
    "if len(azure_openai_key) == 0:\n",
    "    raise Exception(\"API key required\")\n",
    "azure_openai_embedding_model = os.environ[\"AZURE_OPENAI_EMBEDDING_MODEL\"]\n",
    "azure_openai_embedding_deployment = os.environ[\"AZURE_OPENAI_EMBEDDING_DEPLOYMENT\"]\n",
    "azure_openai_chatgpt_deployment = os.environ[\"AZURE_OPENAI_CHATGPT_DEPLOYMENT\"]\n",
    "azure_openai_api_version = os.environ[\"AZURE_OPENAI_API_VERSION\"]\n",
    "embedding_dimensions = int(os.getenv(\"AZURE_OPENAI_EMBEDDING_DIMENSIONS\", 1536))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure an embeddings instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "\n",
    "embeddings = AzureOpenAIEmbedding(\n",
    "    model_name=azure_openai_embedding_model,\n",
    "    deployment_name=azure_openai_embedding_deployment,\n",
    "    api_version=azure_openai_api_version,\n",
    "    azure_endpoint=azure_openai_endpoint,\n",
    "    api_key=azure_openai_key\n",
    ")\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    deployment_name=azure_openai_chatgpt_deployment,\n",
    "    api_version=azure_openai_api_version,\n",
    "    azure_endpoint=azure_openai_endpoint,\n",
    "    api_key=azure_openai_key\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload, vectorize, and index documents\n",
    "\n",
    "This step reads PDFs from a local folder, calls the embedding model for vectorization, and then calls a search client to index the content on Azure AI Search.\n",
    "\n",
    "If you get an exception, such as \"type 'Result' is not subscriptable\", the problem could be an issue with an outdated langchain package. Run `pip show langchain` from a command prompt to get the version, and `pip uninstall` if the package is old."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.azureaisearch import AzureAISearchVectorStore, IndexManagement, MetadataIndexFieldType\n",
    "from llama_index.core import StorageContext, VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.core import Settings\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "\n",
    "metadata_fields = {\n",
    "    \"author\": \"author\",\n",
    "    \"theme\": (\"topic\", MetadataIndexFieldType.STRING),\n",
    "    \"director\": \"director\",\n",
    "}\n",
    "\n",
    "vector_store = AzureAISearchVectorStore(  \n",
    "    search_or_index_client=SearchIndexClient(endpoint=endpoint, credential=credential),  \n",
    "    filterable_metadata_field_keys=metadata_fields,\n",
    "    index_name=index_name,  \n",
    "    index_management=IndexManagement.CREATE_IF_NOT_EXISTS,  \n",
    "    id_field_key=\"id\",  \n",
    "    chunk_field_key=\"content\",  \n",
    "    embedding_field_key=\"content_vector\",  \n",
    "    metadata_string_field_key=\"metadata\",\n",
    "    doc_id_field_key=\"doc_id\",\n",
    "    embedding_dimensionality=embedding_dimensions,\n",
    "    language_analyzer=\"en.lucene\",\n",
    "    vector_algorithm_type=\"exhaustiveKnn\"\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embeddings\n",
    "directory = os.path.abspath(os.path.join(\"..\", \"..\", \"..\", \"..\", \"data\", \"benefitdocs\"))\n",
    "documents = SimpleDirectoryReader(directory).load_data()\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents=documents,\n",
    "    storage_context=storage_context)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform a vector similarity search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Northwind Health Plus plan includes coverage for emergency services, mental health and substance abuse treatment, and out-of-network services, which are not included in the Northwind Standard plan.\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(llm)\n",
    "response = query_engine.query(\"What is included in my Northwind Health Plus plan that is not in standard?\")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
